#include "pch.h"
#include "CloudManager.h"
#include "GaojieDriveApp.h"
#include "AppConfig.h"
#include "PluginManager.h"
#include "RegSyncRoot.h"
#include "DirVisitor.h"
#include "FileReader.h"
#include "..\Common\SRWLock.h"
#include "..\Common\StrHelper.h"
#include "..\Common\Utility.h"

LockExit CCloudManager::m_lockExit;

CCloudManager::CCloudManager()
	: m_connectionKey{}
	, m_hExitEvent(CreateEvent(NULL, TRUE, FALSE, NULL))
{
	AssertInUIThread();

	m_lockExit.reset(m_hExitEvent, SetEvent);
}

CCloudManager::~CCloudManager()
{
	AssertInUIThread();
	assert(!m_lockExit);
	assert(m_mapTasks.empty());

	CloseHandle(m_hExitEvent);
}

bool CCloudManager::Init()
{
	AssertInUIThread();

	CF_SYNC_REGISTRATION registration = { sizeof(registration) };
	registration.ProviderName = L"GaojieDrive";
	registration.ProviderVersion = L"1.0.0.1";

	CF_SYNC_POLICIES policies = { sizeof(policies) };
	policies.Hydration.Primary = CF_HYDRATION_POLICY_FULL;
	policies.Hydration.Modifier = CF_HYDRATION_POLICY_MODIFIER_AUTO_DEHYDRATION_ALLOWED;
	policies.Population.Primary = CF_POPULATION_POLICY_FULL;
	policies.InSync = CF_INSYNC_POLICY_TRACK_FILE_CREATION_TIME | CF_INSYNC_POLICY_TRACK_DIRECTORY_CREATION_TIME;
	policies.HardLink = CF_HARDLINK_POLICY_NONE;

	HRESULT hr = CfRegisterSyncRoot(GetApp()->GetAppConfig()->GetClientRoot().c_str(), &registration, &policies,
		CF_REGISTER_FLAG_UPDATE | CF_REGISTER_FLAG_DISABLE_ON_DEMAND_POPULATION_ON_ROOT);
	assert(hr == S_OK);
	if (hr != S_OK)
	{
		LogError(L"CfRegisterSyncRoot error: hr=0x%x", hr);
		return false;
	}

	AT::AsyncTask(RegisterSyncRoot);

	if (LogInfo())
		InitDeviceMapping();

	CF_CALLBACK_REGISTRATION callbackTable[] =
	{
		{ CF_CALLBACK_TYPE_FETCH_DATA, OnFetchData },
		{ CF_CALLBACK_TYPE_CANCEL_FETCH_DATA, OnCancelFetchData },
		{ CF_CALLBACK_TYPE_FETCH_PLACEHOLDERS, OnFetchPlaceholders },
		{ CF_CALLBACK_TYPE_CANCEL_FETCH_PLACEHOLDERS, OnCancelFetchPlaceholders },
		CF_CALLBACK_REGISTRATION_END
	};

	hr = CfConnectSyncRoot(GetApp()->GetAppConfig()->GetClientRoot().c_str(), callbackTable, NULL,
		CF_CONNECT_FLAG_REQUIRE_PROCESS_INFO | CF_CONNECT_FLAG_REQUIRE_FULL_FILE_PATH | CF_CONNECT_FLAG_BLOCK_SELF_IMPLICIT_HYDRATION, &m_connectionKey);
	assert(hr == S_OK);
	if (hr != S_OK)
	{
		LogError(L"CfConnectSyncRoot error: hr=0x%x", hr);
		return false;
	}

	AT::AsyncTask(InitPlaceholders);
	return true;
}

void CCloudManager::Exit()
{
	AssertInUIThread();

	// Notify and wait for all tasks to exit
	CancelTasks();
	std::atomic_store(&m_lockExit, {});
	WaitForSingleObject(m_hExitEvent, INFINITE);

	HRESULT hr = CfDisconnectSyncRoot(m_connectionKey);
	assert(hr == S_OK);
	if (hr != S_OK)
		LogWarn(L"CfDisconnectSyncRoot error: hr=0x%x", hr);
/*
	hr = CfUnregisterSyncRoot(GetApp()->GetAppConfig()->GetClientRoot().c_str());
	assert(hr == S_OK);
	if (hr != S_OK)
		LogWarn(L"CfUnregisterSyncRoot error: hr=0x%x", hr);
*/
}

void CCloudManager::AddTask(const std::shared_ptr<ICfTask>& pTask)
{
	AssertInThreadpool();

	CExclusiveLock lock(m_SRWLock);
	m_mapTasks[pTask->GetKey()] = pTask;
}

void CCloudManager::RemoveTask(INT64 nTaskKey, bool bCancel)
{
	AssertInThreadpool();

	std::shared_ptr<ICfTask> pTask;
	CExclusiveLock lock(m_SRWLock);

	if (auto it = m_mapTasks.find(nTaskKey); it != m_mapTasks.end())
	{
		pTask = it->second.lock();
		m_mapTasks.erase(it);
	}

	lock.Unlock();

	if (pTask && bCancel)
		pTask->Cancel();
}

void CCloudManager::CancelTasks()
{
	AssertInUIThread();

	CExclusiveLock lock(m_SRWLock);
	auto mapTasks = std::move(m_mapTasks);
	lock.Unlock();

	for (const auto& [key, wpTask] : mapTasks)
	{
		if (auto pTask = wpTask.lock())
			pTask->Cancel();
	}
}

void CCloudManager::InitDeviceMapping()
{
	wchar_t szName[4] = { 'A', ':' }, szPath[256];

	for (DWORD dwDrives = GetLogicalDrives(); dwDrives && *szName <= 'Z'; dwDrives >>= 1, ++*szName)
	{
		if ((dwDrives & 1) && QueryDosDeviceW(szName, szPath, ARRAYSIZE(szPath)))
			m_vecDosDevices.emplace_back(*szName, szPath);
	}
}

void CCloudManager::LogCallbackInfo(LPCWSTR lpType, const CF_CALLBACK_INFO *pCallbackInfo)
{
	LPCWSTR lpNormalizedPath = pCallbackInfo->NormalizedPath;
	LPCWSTR lpProcessPath = pCallbackInfo->ProcessInfo ? pCallbackInfo->ProcessInfo->ImagePath : nullptr;

	// Remove an extra backslash
	if (lpNormalizedPath && *lpNormalizedPath == '\\')
	{
		while (lpNormalizedPath[1] == '\\')
			lpNormalizedPath++;
	}

	for (const auto& it : m_vecDosDevices)
	{
		size_t nNameLen = it.second.size();

		if (lpProcessPath && _wcsnicmp(lpProcessPath, it.second.c_str(), nNameLen) == 0 && lpProcessPath[nNameLen] == '\\')
		{
			LogInfo(L"%s: Path=%s%s, Process=%c:%s", lpType, pCallbackInfo->VolumeDosName, lpNormalizedPath, it.first, lpProcessPath + nNameLen);
			return;
		}
	}

	LogInfo(L"%s: Path=%s%s, Process=%s", lpType, pCallbackInfo->VolumeDosName, lpNormalizedPath, lpProcessPath);
}

// Initialize placeholders for all subdirs
void CCloudManager::InitPlaceholders()
{
	AssertInThreadpool();

	std::vector<std::wstring> vecFileNames;

	// TODO: Support subdirectories
	for (const auto& subDirCfg : GetApp()->GetAppConfig()->GetSubDirCfgs())
	{
		std::wstring strSubPath = subDirCfg.strSubPath;

		if (auto nPos = strSubPath.find('\\'); nPos != std::wstring::npos)
			strSubPath.erase(nPos);

		vecFileNames.emplace_back(std::move(strSubPath));
	}

	CreatePlaceholders(GetApp()->GetAppConfig()->GetClientRoot(), vecFileNames);

	// TODO: Recursively change all subdirectories to ON_DEMAND
}

void CCloudManager::CreatePlaceholders(const std::wstring& strPathName, const std::vector<std::wstring>& vecFileNames)
{
	AssertInThreadpool();

	std::unordered_map<std::wstring, FIND_DATA> mapLocalInfos;

	Util::EnumDirectory((strPathName + L'*').c_str(), [&mapLocalInfos](FIND_DATA&& findData)
	{
		mapLocalInfos.emplace(Str::StringToLower(findData.cFileName), std::move(findData));
		return true;
	});

	std::vector<FileInfo> vecFileInfos;

	for (const auto& strFileName : vecFileNames)
	{
		auto it = mapLocalInfos.find(Str::StringToLower(strFileName));
		if (it != mapLocalInfos.end())
		{
			const auto& findData = it->second;

			// Local directory already exists, just update
			if (findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
			{
				CF::UpdatePlaceholder(strPathName, findData, FileInfo{ findData.cFileName, 0, findData.ftLastWriteTime, findData.dwFileAttributes });
				continue;
			}

			// Delete existing files with the same name
			std::wstring strFileName = strPathName + findData.cFileName;

			if (findData.dwFileAttributes & FILE_ATTRIBUTE_READONLY)
				SetFileAttributesW(strFileName.c_str(), FILE_ATTRIBUTE_NORMAL);

			DeleteFileW(strFileName.c_str());
		}

		vecFileInfos.emplace_back(FileInfo{ strFileName, 0, {}, FILE_ATTRIBUTE_DIRECTORY });
	}

	// Create placeholders that don't exist locally
	if (vecFileInfos.size())
		CF::CreatePlaceholders(strPathName, vecFileInfos);
}

// Get subpath relative to ClientRoot
LPCWSTR CCloudManager::GetSubPathOnRoot(LPCWSTR lpSubPath)
{
	if (!lpSubPath)
		return lpSubPath;

	while (*lpSubPath == '\\')
		lpSubPath++;

	// PathSkipRoot
	std::wstring_view strClientRoot = GetApp()->GetAppConfig()->GetClientRoot();
	strClientRoot.remove_prefix(strClientRoot.find(':') + 2);

	if (_wcsnicmp(lpSubPath, strClientRoot.data(), strClientRoot.size()) == 0)
		lpSubPath += strClientRoot.size();
	else
		lpSubPath = nullptr;

	return lpSubPath;
}

void CCloudManager::OnFetchData(const CF_CALLBACK_INFO *pCallbackInfo, const CF_CALLBACK_PARAMETERS *pCallbackParameters)
{
	AssertInThreadpool();

	auto lockExit = std::atomic_load(&m_lockExit);
	if (!lockExit)
		return;

	if (LogInfo())
		GetApp()->GetCloudManager()->LogCallbackInfo(L"OnFetchData", pCallbackInfo);

	LPCWSTR lpSubPath = GetSubPathOnRoot(pCallbackInfo->NormalizedPath);
	if (lpSubPath && *lpSubPath)
	{
		IPluginPtr pPlugin;
		size_t nSubLen = GetApp()->GetPluginManager()->GetPlugin(lpSubPath, pPlugin);

		if (pPlugin)
		{
			auto pReader = std::make_shared<CFileReader>(pCallbackInfo, std::move(lockExit), GetApp()->GetAppConfig()->GetClientRoot() + lpSubPath);
			GetApp()->GetCloudManager()->AddTask(pReader);
			pPlugin->ReadFile(lpSubPath + nSubLen, pReader, pCallbackParameters->FetchData.RequiredFileOffset.QuadPart, pCallbackParameters->FetchData.RequiredLength.QuadPart);
			return;
		}
	}

	assert(false);
	LogError(L"OnFetchData error: %s", pCallbackInfo->NormalizedPath);

	CF::CCbWrapper cfCbWrapper(pCallbackInfo);
	cfCbWrapper.TransferData(pCallbackInfo->NormalizedPath, nullptr, pCallbackParameters->FetchData.RequiredFileOffset.QuadPart, pCallbackParameters->FetchData.RequiredLength.QuadPart);
}

void CCloudManager::OnCancelFetchData(const CF_CALLBACK_INFO *pCallbackInfo, const CF_CALLBACK_PARAMETERS *pCallbackParameters)
{
	AssertInThreadpool();

	auto lockExit = std::atomic_load(&m_lockExit);
	if (!lockExit)
		return;

	GetApp()->GetCloudManager()->RemoveTask(pCallbackInfo->TransferKey.QuadPart, true);
}

void CCloudManager::OnFetchPlaceholders(const CF_CALLBACK_INFO *pCallbackInfo, const CF_CALLBACK_PARAMETERS *pCallbackParameters)
{
	AssertInThreadpool();

	auto lockExit = std::atomic_load(&m_lockExit);
	if (!lockExit)
		return;

	if (LogInfo())
		GetApp()->GetCloudManager()->LogCallbackInfo(L"OnFetchPlaceholders", pCallbackInfo);

	LPCWSTR lpSubPath = GetSubPathOnRoot(pCallbackInfo->NormalizedPath);
	if (lpSubPath && *lpSubPath)
	{
		std::wstring strSubPath = lpSubPath;
		Util::FormatFilePath(strSubPath);

		IPluginPtr pPlugin;
		size_t nSubLen = GetApp()->GetPluginManager()->GetPlugin(strSubPath, pPlugin);

		if (pPlugin)
		{
			auto pVisitor = std::make_shared<CDirVisitor>(pCallbackInfo, std::move(lockExit), GetApp()->GetAppConfig()->GetClientRoot() + strSubPath);
			GetApp()->GetCloudManager()->AddTask(pVisitor);
			pVisitor->Init();
			pPlugin->VisitDir(strSubPath.c_str() + nSubLen, pVisitor);
			return;
		}
	}

	assert(false);
	LogError(L"OnFetchPlaceholders error: %s", pCallbackInfo->NormalizedPath);

	CF::CCbWrapper cfCbWrapper(pCallbackInfo);
	cfCbWrapper.TransferPlaceholders(pCallbackInfo->NormalizedPath, true, {});
}

void CCloudManager::OnCancelFetchPlaceholders(const CF_CALLBACK_INFO *pCallbackInfo, const CF_CALLBACK_PARAMETERS *pCallbackParameters)
{
	AssertInThreadpool();

	auto lockExit = std::atomic_load(&m_lockExit);
	if (!lockExit)
		return;

	GetApp()->GetCloudManager()->RemoveTask(pCallbackInfo->TransferKey.QuadPart, true);
}
